package com.example;

import org.jnetpcap.*;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Random;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

public class Socket {
    public int srcPort;
    public int desPort;
    public byte[] desIP;

    // TCP序列号和确认号
    /**
     * seq是我方发送的下一个数据包的序列号
     * 每次发送完一个数据包后，seq += 数据部分的长度
     * 如果发送的数据包是SYN和FIN，还需要额外+1
     */
    int seq;
    /**
     * 这个ack是我方期望接收到的下一个数据包的序列号，
     * 在接收到对方的数据包后，先检查是否是期望的数据包，如果是，更新ack为对方的seq + 数据部分的长度，
     * 如果收到的数据包是SYN和FIN，还需要额外+1。<br>
     * 如果不是期望的数据包，不更新ack。todo 此处的处理可以加入快速重传机制
     */
    int ack;

    /**
     * 对方的窗口大小，默认为65535
     */
    int windowSize;

    /**
     * 用于控制发送线程何时退出
     */
    AtomicBoolean running = new AtomicBoolean(true);

    PcapIf pcapIf;

    Socket(byte[] desIP, int srcPort, int desPort) {
        this.desIP = desIP;
        this.desPort = desPort;
        this.srcPort = srcPort;

        Random rand = new Random();
        //generate random values from 1-100000
        seq = rand.nextInt(99999) + 1;
        ack = 0;
        windowSize = 65535;

        try {
            pcapIf = Utils.getPcapIf();
        } catch (Exception e) {
            throw new RuntimeException("No device found", e);
        }

    }

    public static ClosedSocket newSocket(byte[] desIP, int desPort) {
        return new ClosedSocket(desIP, desPort);
    }

    // 用于发送SYN包，接收SYN-ACK包，发送ACK包，完成三次握手
    // 三次握手完成后，返回一个OpenSocket对象，用于发送和接收数据
    // todo CloseSocket没有使用PacketBufferQueue作为缓冲区，写法也比较冗余，需要重构
    public static class ClosedSocket extends Socket {

        final ByteBuffer sendBuffer;
        final ByteBuffer recvBuffer;

        ClosedSocket(byte[] desIP, int desPort) {
            this(desIP, generateRandomPort(), desPort);
        }

        ClosedSocket(byte[] desIP, int srcPort, int desPort) {
            super(desIP, srcPort, desPort);
            // 初始化缓冲区，容量为1518 * 100字节，即100个以太帧的大小
            sendBuffer = ByteBuffer.allocate(1518 * 100);
            sendBuffer.limit(0);

            recvBuffer = ByteBuffer.allocate(1518 * 100);
            recvBuffer.limit(0);
        }

        private static int generateRandomPort() {
            Random rand = new Random();
            int lowerBound = 1024;
            int upperBound = 65535;
            //generate random values from 1024-65535
            int srcPort = rand.nextInt(upperBound - lowerBound) + lowerBound;
//            System.out.printf("Random TCP port value from %d to %d : %d%n", lowerBound, upperBound, srcPort);
            return srcPort;
        }

        public OpenSocket connect() {
            InetAddress desIPAddress;
            try {
                desIPAddress = InetAddress.getByAddress(desIP);
            } catch (UnknownHostException e) {
                throw new RuntimeException(e);
            }

            // 开启两个线程，一个用于接收数据，一个用于发送数据
            // 接收数据的线程只需要在接收到数据后，将数据存到缓冲区中即可。
            CompletableFuture.runAsync(() -> {
                try (Pcap pcap = Pcap.openLive(pcapIf, 64 * 1024, false, 5 * 1000, TimeUnit.MILLISECONDS);) {

//                    System.out.println("打开网络接口1");
                    String filter = "src host %s and dst port %s and src port %s".formatted(desIPAddress.getHostAddress(), srcPort, desPort);
//                    System.out.println("Filter: " + filter);
                    BpFilter bpFilter = pcap.compile(filter, false);
                    pcap.setFilter(bpFilter);
//                    pcap.setNonBlock(true);

                    pcap.loop(1, this::receiverHandler, pcap);
//                    System.out.println("开始接收数据...");

                } catch (PcapException e) {
                    throw new RuntimeException(e);
                }
            });
            System.out.printf("Connecting to %s:%d...\n", desIPAddress.getHostAddress(), desPort);


            // 睡眠三秒等待接收线程启动
            try {
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }

            Semaphore semaphore = new Semaphore(0);

            // 发送数据的线程只需要在缓冲区有数据时，将数据发送出去。
            CompletableFuture.runAsync(() -> {
                try (Pcap pcap = Pcap.create(pcapIf)) {
                    pcap.setSnaplen(16 * 1024);
                    pcap.activate();

//                    System.out.println("打开网络接口2");

                    while (running.get()) {
                        synchronized (sendBuffer) {
                            if (sendBuffer.hasRemaining()) {
//                                System.out.println();
                                try {
//                                    byte[] packet = sendBuffer.slice().array();
//                                    pcap.sendPacket(packet);
                                    pcap.sendPacket(Arrays.copyOf(sendBuffer.array(), sendBuffer.limit()));
                                    sendBuffer.position(0).limit(0); // 清空缓冲区
                                    // 发送完一个包后，释放一个许可
                                    semaphore.release();
                                } catch (PcapException e) {
                                    throw new RuntimeException("sending packet failed", e);
                                }
                            } else {
                                System.out.print(".");
                            }
                        }
                        try {
                            Thread.sleep(100); // 每100毫秒检查一次
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                    }

                } catch (PcapException e) {
                    throw new RuntimeException(e);
                }

            });

            // 主线程，往缓冲区中放入SYN包，等待接收SYN-ACK包，然后发送ACK包。
            var SYNpacket = Packet.builder(Packet.PktType.SYN, desIP, (short) srcPort, (short) desPort, seq, ack).build();
            seq += 1;
            synchronized (sendBuffer) {
                sendBuffer.limit(SYNpacket.bytes().length);
                sendBuffer.put(SYNpacket.bytes());
                sendBuffer.position(0);
            }

            // 使用semaphore.acquire()方法来等待数据包发送成功的信号
            // 如果发送数据包的操作还没有完成，acquire()方法会阻塞直到操作完成
            try {
                semaphore.acquire();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

//            System.out.println(recvBuffer.remaining());

            byte[] receivedData = pollReceiveBuffer();

            Packet recvPacket = Packet.buildFromBytes(receivedData);
//            System.out.println("Received packet");

            // 如果收到的数据包是SYN-ACK包，发送ACK包
            // todo 暂时不考虑其他情况，如同时握手等
            if (recvPacket.isSYN() && recvPacket.isACK()) {
//                System.out.println("SYN-ACK packet captured");

                windowSize = recvPacket.windowSize();

                ack = recvPacket.seqNumber() + 1;
                var ackPacket = Packet.builder(Packet.PktType.ACK, desIP, (short) srcPort, (short) desPort, seq, ack).build();
                byte[] ackPacketBytes = ackPacket.bytes();
                synchronized (sendBuffer) {
                    sendBuffer.limit(ackPacketBytes.length);
                    sendBuffer.put(ackPacketBytes);
                    sendBuffer.position(0);
                }
//                System.out.println("ACK packet sent");
            }
            // 主线程等待子线程发包成功
            try {
                semaphore.acquire();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            running.set(false);
            System.out.println("\nConnection established\n");
            return new OpenSocket(this);
        }

        private void receiverHandler(Pcap pcap, PcapHeader header, byte[] packet) {
//            System.out.println("\nReady to receive packet");
            synchronized (recvBuffer) {
//                System.out.println("Receiving packet");
                recvBuffer.limit(packet.length);
                recvBuffer.put(packet);
                recvBuffer.flip();
                recvBuffer.notifyAll();
            }
        }

        private byte[] pollReceiveBuffer() {
            byte[] data = null;
            synchronized (recvBuffer) {
                while (!recvBuffer.hasRemaining()) {
                    try {
                        recvBuffer.wait();
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt(); // Restore the interrupted status
                        throw new RuntimeException(e);
                    }
                }
                data = new byte[recvBuffer.remaining()];
                recvBuffer.get(data);
                recvBuffer.clear();
            }
            return data;
        }
    }

    public static class OpenSocket extends Socket {

        final PacketBufferQueue sendBuffer;
        final PacketBufferQueue recvBuffer;

        CompletableFuture<Void> receiver;
        CompletableFuture<Void> sender;

        boolean FIN_WAIT = false;

        int temp = 0;

        OpenSocket(ClosedSocket closedSocket) {
            super(closedSocket.desIP, closedSocket.srcPort, closedSocket.desPort);
            this.ack = closedSocket.ack;
            this.seq = closedSocket.seq;
            this.windowSize = closedSocket.windowSize;

            sendBuffer = new PacketBufferQueue(1518 * 100, ack);
            recvBuffer = new PacketBufferQueue(1518 * 100);


            // 开启两个线程，一个用于接收数据，一个用于发送数据
            // 接收数据的线程只需要在接收到数据后，将数据存到缓冲区中即可。
            receiver = CompletableFuture.runAsync(() -> {
                try (Pcap pcap = Pcap.openLive(pcapIf, 64 * 1024, false, 5 * 1000, TimeUnit.MILLISECONDS);) {

//                    System.out.println("打开接收缓冲区");
                    InetAddress address = InetAddress.getByAddress(desIP);
                    String filter = "src host %s and dst port %s and src port %s".formatted(address.getHostAddress(), srcPort, desPort);
//                    System.out.printf("Filter: %s%n", filter);
                    BpFilter bpFilter = pcap.compile(filter, false);
                    pcap.setFilter(bpFilter);

                    // 无限循环，直到收到FIN包
                    pcap.loop(-1, this::receiverHandler, pcap);
                } catch (PcapException | UnknownHostException e) {
                    throw new RuntimeException(e);
                }
            });

            // 睡眠三秒等待接收线程启动
            try {
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }

//            Semaphore semaphore = new Semaphore(0);

            // 发送数据的线程只需要在缓冲区有数据时，将数据发送出去。
            sender = CompletableFuture.runAsync(() -> {
                try (Pcap pcap = Pcap.create(pcapIf)) {
                    pcap.setSnaplen(16 * 1024);
                    pcap.activate();

//                    System.out.println("打开发送缓冲区");

                    while (running.get() || !sendBuffer.isEmpty()) {
                        synchronized (sendBuffer) {
                            while (sendBuffer.isEmpty() && running.get()) {
                                try {
//                                    System.out.println("No data to send, waiting...");
                                    sendBuffer.wait(); // 等待，直到其他线程调用sendBuffer.notify()
                                } catch (InterruptedException e) {
                                    Thread.currentThread().interrupt();
                                    throw new RuntimeException(e);
                                }
                            }

                            if (!sendBuffer.isEmpty()) {
                                try {
                                    pcap.sendPacket(sendBuffer.readPacket());
                                } catch (PcapException e) {
                                    throw new RuntimeException("sending packet failed", e);
                                }
                            }
                        }
                    }
                } catch (PcapException e) {
                    throw new RuntimeException(e);
                }
            });

        }

        // send() 方法可以发送超过一个MTU的数据，会自动分片。
        public void send(byte[] data) {
            if (FIN_WAIT) {
                throw new RuntimeException("Socket is in FIN_WAIT state, cannot send data");
            }

            if (data.length > windowSize) {
                // 仅发送在对方窗口大小内的数据
                // 其实应该等一会再发送剩下的数据，但是这里简化处理，直接截断数据
                data = Arrays.copyOf(data, windowSize);
            }

            // 获取我方窗口大小
            int myWindowSize;
            synchronized (recvBuffer) {
                myWindowSize = recvBuffer.unusedSize();
            }

            // 分片发送数据
            int numOfPackets = data.length / 1460 + 1;

            for (int i = 0; i < numOfPackets; i++) {
                int start = i * 1460;
                int end = Math.min((i + 1) * 1460, data.length);
                byte[] dataSlice = Arrays.copyOfRange(data, start, end);
                // 对于没有设置ACK标志位的包，ackNumber为0
                Packet packet = Packet.builder(dataSlice, desIP, (short) srcPort, (short) desPort, seq, ack)
                        .setACK()
                        .setWindowSize(myWindowSize)
                        .build();
                synchronized (sendBuffer) {
                    sendBuffer.putPacket(packet.bytes());
                    seq += dataSlice.length; // 更新序列号
                    sendBuffer.want(seq);
                    sendBuffer.notify();
                }
            }
        }

        // 调用此方法会发送一个FIN包，代表我方不再发送新的数据包，但是仍能接收新的数据包直到收到对方的FIN。
        // 理论上应该将此对象转换成一个新类型（状态），新的类没有send()方法，这样能完全避免close()后用户仍然尝试发送新数据
        // 但是这里简化处理，直接在close()方法中设置一个标志位，send()方法中检查这个标志位，如果为true，不再发送数据。
        public void close() {

            System.out.println("Closing connection...");


            int myWindowSize;
            synchronized (recvBuffer) {
                myWindowSize = recvBuffer.unusedSize();
            }
//            seq += 1;
            Packet finPacket = Packet.builder(Packet.PktType.FIN, desIP, (short) srcPort, (short) desPort, seq, ack)
                    .setWindowSize(myWindowSize)
                    .setACK()
                    .build();
            synchronized (sendBuffer) {
                sendBuffer.putPacket(finPacket.bytes());
                seq += 1;
                sendBuffer.want(seq);
                sendBuffer.notify();
            }

            FIN_WAIT = true;

            // 阻塞调用close()方法的线程，直到本Socket的发送线程与接收线程都结束
            CompletableFuture.allOf(sender, receiver).join();
        }

        // 如果数据包有多个分片，只返回第一个分片的数据。需要在外部循环调用，直到返回null为止。
        // 如果没有数据包，会阻塞直到有数据包。
        // 如果阻塞超过7秒，返回null。
        public byte[] recv() {
            var selector = pollRecvBuffer();
            try {
                return selector.get(10, TimeUnit.SECONDS); // todo 设定为更合理的超时时间
            } catch (TimeoutException e) {
                selector.cancel(true);
                return null;
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(e);
            }
        }

        // 用于接收数据包的回调函数
        // 传入的Pcap对象是用于接收数据包的Pcap对象，用于在接收到FIN包后停止抓包循环，关闭接收线程。
        private void receiverHandler(Pcap pcap, PcapHeader header, byte[] packetBytes) {
            Packet recvPacket = Packet.buildFromBytes(packetBytes);

//            temp += 1;

            if (recvPacket.isACK()) {
//                System.out.printf("No.%d: Received ACK packet. ack: %d. seq: %d%n", temp, recvPacket.ackNumber(), recvPacket.seqNumber());
                int ackNumber = recvPacket.ackNumber();
                synchronized (sendBuffer) {
                    sendBuffer.compact(ackNumber);
                }
            }

            // 在不考虑快速重传的情况下，如果此数据包不是FIN、不是SYN、也没有数据，对于本程序来说，认为它只是一个纯ACK包。
            // 此时不需要发送任何回复。直接返回
            if (!recvPacket.isFIN() && recvPacket.data().length == 0) {
                return;
            }

            // 其他任何情况都需要发一个ACK包作为回复。
            var packetBuilder = Packet.builder(Packet.PktType.ACK, desIP, (short) srcPort, (short) desPort, seq, ack);

            // 如果此包仅带有标志位，不带数据，则没有必要放入缓冲区
            // 这也意味着如果收到一个没有数据、也没有标志位的包，那么这个包会被直接丢弃
            // 只将数据包携带的数据放入缓冲区
            if (recvPacket.data().length > 0) {
//                System.out.println("Receiving data packet");
                int unusedBufferSize;
                synchronized (recvBuffer) {
                    recvBuffer.putPacket(packetBytes);
                    unusedBufferSize = recvBuffer.unusedSize();
                    recvBuffer.notifyAll();
                }
//                System.out.printf("No.%d: Received data packet. data length: %d, ack: %d, seq: %d%n", temp, recvPacket.data().length, recvPacket.ackNumber(), recvPacket.seqNumber());


                // 更新ack
                ack = recvPacket.seqNumber() == ack ? recvPacket.seqNumber() + recvPacket.data().length : ack; // 快速重传需要的部分特性
                // 由于此ACK包不包含数据，因此不需要更新seq

                // 更新对方的窗口大小
                windowSize = recvPacket.windowSize();

//                packetBuilder.setAckNumber(ack);
                packetBuilder.setAckNumber(ack).setWindowSize(unusedBufferSize);

            }

            if (recvPacket.isFIN()) {
                ack += 1;
                packetBuilder.setAckNumber(ack);
            }

            Packet ACK = packetBuilder.build();

            // 收到数据包后，回复ACK
            synchronized (sendBuffer) {
                sendBuffer.putPacket(ACK.bytes());
                sendBuffer.notify();

            }
//            System.out.printf("No.%d: Sending ACK packet. ack: %d. seq: %d%n", temp, ACK.ackNumber(), ACK.seqNumber());


            // 假设服务器不会主动关闭连接，只有我方会主动关闭连接，
            // 因此如果收到FIN包，一定是我方已经发了FIN包过去，对方回复了FIN包，
            // 此时双方都不再发送数据包。回复一个ACK包后，可以关掉缓冲区。
            if (recvPacket.isFIN()) {
                // 理论上此时应该进入TIME_WAIT状态，等待2MSL时间没有收到新包后关闭连接，如果还能收到新包，需要重传FIN
                // 但这2MSL等于4分钟（RFC793定义MSL为2分钟），太久了，这里简化处理，直接关闭连接

                running.set(false); // 停止发送缓冲区的轮询，关闭发送线程
                pcap.breakloop(); // 停止抓包循环，关闭接收线程
//                System.out.println("No." + temp + ": FIN packet captured");
            }
        }

        // 似乎没有必要使用CompletableFuture，但如果未来要扩展超时重传等功能，可以参考这个函数的写法。
        private CompletableFuture<byte[]> pollRecvBuffer() {
            return CompletableFuture.supplyAsync(() -> {
                byte[] data = null;
                synchronized (recvBuffer) {
                    while (recvBuffer.isEmpty()) {
                        try {
                            // 为了避免过于频繁的轮询，这里使用wait()方法。
                            // 每次往接收缓冲区放入数据后，会调用notifyAll()方法唤醒所有等待的线程
                            recvBuffer.wait();
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt(); // Restore the interrupted status
                            throw new RuntimeException(e);
                        }
                    }
                    data = recvBuffer.readPacket();
                    recvBuffer.compact();
                }
                return data;
            });
        }
    }

    /**
     * 用于存储数据包的缓冲区。<br>
     * 这是一个先进先出的队列，用于存储数据包，每次put和read都只操作一个数据包。
     */
    public static class PacketBufferQueue {
        // 由于对缓冲区的轮询通过检查hasRemaining()来判断是否为空，因此需要手动维护position和limit
        //
        // buffer有两个状态，写入状态和读取状态。
        // 读取操作远比写入操作频繁，因此默认状态是待读取状态，切换成写入状态并写入完成后需要切换回待读取状态。
        //
        // 读取过程如下：
        //  待读取状态：
        //    +-------+-----------------------------+---------------------+
        //    |       |            Data             |    Unused Space     |
        //    +-------+-----------------------------+---------------------+
        //    ^       ^                             ^                     ^
        //    |       |                             |                     |
        //    0    position                       limit               capacity
        //
        //  读取完成：
        //    +-------------+-----------------------+---------------------+
        //    |  Read Data  |         Data          |    Unused Space     |
        //    +-------------+-----------------------+---------------------+
        //    ^             ^                       ^                     ^
        //    |             |                       |                     |
        //    0          position                 limit               capacity
        //
        //
        // 写入过程如下：
        //  先进行标记，然后移动位置，将待读取状态变为待写入状态：
        //    +-----+-------------+---------------------------------------+
        //    |     |    Data     |             Unused Space              |
        //    +-----+-------------+---------------------------------------+
        //    ^     ^             ^                                       ^
        //    |     |             |                                       |
        //    0    mark       position                            capacity = limit
        //
        //  写入完成：
        //    +-----+-------------+----------+----------------------------+
        //    |     |    Data     | New Data |       Unused Space         |
        //    +-----+-------------+----------+----------------------------+
        //    ^     ^                        ^                            ^
        //    |     |                        |                            |
        //    0    mark                  position                 capacity = limit
        //
        //  复原位置：
        //    +-----+-------------+----------+----------------------------+
        //    |     |    Data     | New Data |       Unused Space         |
        //    +-----+-------------+----------+----------------------------+
        //    ^     ^                        ^                            ^
        //    |     |                        |                            |
        //    0   position                 limit                       capacity
        //
        //
        // 清除已读取或已收到ACK的数据：
        //  将未读取的数据移到缓冲区头部，然后将position置零，limit置为未读取数据的长度
        //    +----------------+------------------------------------------+
        //    |      Data      |                Unused Space              |
        //    +----------------+------------------------------------------+
        //    ^                ^                                          ^
        //    |                |                                          |
        //  position = 0     limit                                   capacity
        //
        private final ByteBuffer buffer;

        // 对于发送缓冲区：
        // 发送数据包后，记录下期望的对应此数据包的ACK，当收到对方确认后，清除已收到ACK确认的数据包
        //
        // 对于接收缓冲区，不需要记录这个值，
        // 因为收到数据包后我方立即给对方回复ACK，然后当用户程序读取数据包时，直接清除已读取的数据包
        private int wantingAck = -1;

        public PacketBufferQueue(int capacity, int ack) {
            this.buffer = ByteBuffer.allocate(capacity);
            this.buffer.limit(0);
            this.wantingAck = ack;
        }

        public PacketBufferQueue(int capacity) {
            this.buffer = ByteBuffer.allocate(capacity);
            this.buffer.limit(0);
        }

        /**
         * 仅适用于发送缓冲区。<br>
         * 当发送数据包后，记录下一个应该发送的数据包的seq作为期望的ack。
         * 如果收到了对方的ACK包，且ack等于期望的ack，说明所有已发送的数据包都已被接收，可以清除缓冲区中已发送的数据包。
         *
         * @param newWanting 下一个应该发送的数据包的seq作为期望的ack
         */
        public void want(int newWanting) {
            if (newWanting > wantingAck) {
                wantingAck = newWanting;
            }
        }

        public int wanting() {
            if (wantingAck == -1) {
                throw new RuntimeException("No wanting ack, wanting ack should only be used for sending buffer");
            }
            return wantingAck;
        }

        public void putPacket(byte[] packet) {
            buffer.mark().position(buffer.limit()).limit(buffer.capacity());
            if (buffer.remaining() < packet.length + 4) {
                throw new RuntimeException("Buffer is full");
            }
            buffer.putInt(packet.length);
            buffer.put(packet);
            buffer.limit(buffer.position()).reset();
        }

        // 需要在外部检查是否有数据，确定有数据再调用本方法
        public byte[] readPacket() {
            if (buffer.remaining() < 4) {
                throw new RuntimeException("Buffer is empty");
            }
            int packetLength = buffer.getInt();
            byte[] packetBytes = new byte[packetLength];
            buffer.get(packetBytes);

//            int remaining = buffer.remaining();
//            buffer.compact().limit(remaining).position(0);

            return packetBytes;
        }

        public void compact() {
            int remaining = buffer.remaining();
            buffer.compact().limit(remaining).position(0);
        }

        /**
         * 仅适用于发送缓冲区。<br>
         * 当收到对方的ACK后，如果ack等于期望的ack，清除已发送的数据包
         *
         * @param newAck 收到的ack
         */
        public void compact(int newAck) {
            if (newAck == wantingAck) {
                compact();
            }
        }

        public int unusedSize() {
            return buffer.capacity() - buffer.limit();
        }

        public boolean isEmpty() {
            return !buffer.hasRemaining();
        }
    }
}